import cv2 as cv
import numpy as np


def GassianPyramidDemo(img, n_level):
    cv.imshow(f'GauPyr#0', img)
    gasPyr = []
    tmp = img
    for level_id in range(n_level):
        down = cv.pyrDown(tmp)
        cv.imshow(f'GauPyr#{level_id+1}', down)
        gasPyr.append(down)
        tmp = down
    return gasPyr


def LaplacianPyramidDemo(img, n_level):
    gasPyr = GassianPyramidDemo(img, n_level)
    gasPyr.insert(0, img)
    lapPyr = []
    for level_id in range(n_level):
        up_level_id = level_id + 1
        up = cv.pyrUp(gasPyr[up_level_id])
        gas = gasPyr[level_id]
        lap = cv.subtract(gas, up)
        cv.imshow(f'LapPyr#{level_id+1}', lap)
        lapPyr.append(lap)
    return lapPyr


if '__main__' == __name__:
    # path = '../../../../large_data/pic/len_top.jpg'
    path = '../../../../large_data/pic/portrait/baby.jpg'
    # path = '../../../../large_data/pic/watershed/water_coins.jpg'
    # path = '../../../../large_data/pic/weapon/f22.mft-5.jpg'  # no effect because of small pic => 512x512
    img = cv.imread(path, cv.IMREAD_COLOR)
    cv.imshow('ori', img)
    H, W, CH = img.shape[:3]
    print(H, W, CH)
    if H < W:
        padding = W - H
        img = np.concatenate([img, np.zeros([padding, W, CH], dtype=np.uint8)], axis=0)
    else:
        padding = H - W
        img = np.concatenate([img, np.zeros([H, padding, CH], dtype=np.uint8)], axis=1)
    img = cv.resize(img, (512, 512), interpolation=cv.INTER_CUBIC)
    H, W, CH = img.shape[:3]
    print(H, W, CH)

    LaplacianPyramidDemo(img, 3)

    cv.waitKey(0)
    cv.destroyAllWindows()
